package com.warmheart.base.executor;

import java.util.Date;

import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;

import com.warmheart.core.util.UUIDGenerator;

public abstract class TraceContext {

    private static final ThreadLocal<Date> timeThreadLocal = new ThreadLocal<Date>();

    public static void start() {
        Date startTime = new Date();
        timeThreadLocal.set(startTime);
    }

    public static long end() {
        Date endTime = new Date();
        Date startTime = timeThreadLocal.get();
        if (startTime == null) {
            throw new IllegalArgumentException("必须先调用start方法");
        }
        timeThreadLocal.remove();
        return endTime.getTime() - startTime.getTime();
    }

    public static void init() {
        MDC.clear();
        initTrace();
        initSpan();
    }
    
    public static void clear() {
        timeThreadLocal.remove();
        MDC.clear();
    }

    private static void initTrace() {
        MDC.put(CommonConstants.SLF4J_TRACE_ID, UUIDGenerator.getUUID());
    }

    private static void initSpan() {
        MDC.put(CommonConstants.SLF4j_SPAN_ID, CommonConstants.ROOT_SPAN_ID);
    }

    public static void setSpan() {
        String spanId = MDC.get(CommonConstants.SLF4j_SPAN_ID);
        if (StringUtils.isBlank(spanId)) {
            spanId = CommonConstants.ROOT_SPAN_ID;
        } else {
            spanId = spanId + "." + UUIDGenerator.getUUID();
        }
        MDC.put(CommonConstants.SLF4j_SPAN_ID, spanId);
    }

    public static String getSpan() {
        return MDC.get(CommonConstants.SLF4j_SPAN_ID);
    }

    public static void resetSpan(String spanId) {
        MDC.put(CommonConstants.SLF4j_SPAN_ID, spanId);
    }

    public static Long getSpanCost() {
        String spanCost = MDC.get(CommonConstants.SLF4j_SPAN_COST);
        return Long.valueOf(StringUtils.isNotBlank(spanCost) ? spanCost : "0");
    }

    public static void resetSpanCost(Long costTime) {
        MDC.put(CommonConstants.SLF4j_SPAN_COST, String.valueOf(costTime));
    }
}
